package LibDL.models;

import LibDL.core.*;
import LibDL.core.nn.Module;

import java.util.Arrays;

public class BPR extends Module {
    private Tensor W, H;
    public BPR(int user_size, int item_size, int dim) {
        W = register_parameter("W", functional.rand(new StdVectorLong(Arrays.asList(user_size, dim))));
        H = register_parameter("H", functional.rand(new StdVectorLong(Arrays.asList(item_size, dim))));
    }
/*    public Tensor forward(int u, int i, int j) {
        Tensor row_u = W.index_select(0, ld.tensor(u).to(Dtype.INT64));
        Tensor x_ui = row_u.mul(H.index_select(0, ld.tensor(i).to(Dtype.INT64))).sum(1);
        Tensor x_uj = row_u.mul(H.index_select(0, ld.tensor(j).to(Dtype.INT64))).sum(1);
        Tensor x_uij = x_ui.sub(x_uj);
        Tensor log_prob= ld.log(ld.sigmoid(x_uij)).mean();
        return log_prob.mul(new Scalar(-1));
    }*/
    public Tensor forward(Tensor u, Tensor i, Tensor j) {
        Tensor row_u = W.index_select(0, u);
        Tensor x_ui = row_u.mm(H.index_select(0, i).t()).sum(1);
        Tensor x_uj = row_u.mm(H.index_select(0,j).t()).sum(1);
        Tensor x_uij = x_ui.sub(x_uj);
        Tensor log_prob= functional.log(functional.sigmoid(x_uij)).div(new Scalar(H.size(-1)));
        return log_prob.mul(new Scalar(-1));
    }
/*    public Tensor l2_norm(int u, int i, int j){
        return ld.sum(W.get(u).mul(W.get(u))).add(ld.sum(H.get(i).mul(H.get(i)))).add(ld.sum(H.get(j).mul(H.get(j))));
    }*/
    public Tensor l2_norm(Tensor u, Tensor i, Tensor j){
        Tensor row_u = W.index_select(0, u);
        Tensor row_i = H.index_select(0, i);
        Tensor row_j = H.index_select(0, j);
        Tensor l2u = row_u.mm(row_u.t()).sum(1);
        Tensor l2i = row_i.mm(row_i.t()).sum(1);
        Tensor l2j = row_j.mm(row_j.t()).sum(1);
        return l2u.add(l2i).add(l2j);
    }
    public Tensor predict(int u, Tensor i){
        Tensor row_u = W.index_select(0, functional.tensor(u).to(Dtype.INT64));
        Tensor row_i = H.index_select(0, i);
        Tensor ui = functional.sum(row_i.mul(row_u),-1);
        return  ui;
    }
    public Tensor predict(Tensor u, Tensor i){
        Tensor row_u = W.index_select(0, u);
        Tensor row_i = H.index_select(0, i);
        Tensor ui = row_u.mm(row_i.t());
        return  ui;
    }
}